package example.ThreadTest;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.*;

public class MultiThreadNioServer {
    private Charset charset = Charset.forName("UTF-8");
    private ExecutorService bossExecutor;    // 主 Reactor 线程池
    private ExecutorService workerExecutor;  // 从 Reactor 线程池
    private ExecutorService businessExecutor;// 业务线程池
    
    // 从 Reactor 组
    private List<SubReactor> subReactors = new ArrayList<>();
    private int subReactorIndex = 0;
    
    public MultiThreadNioServer() throws IOException {
        // 主 Reactor：1个线程处理连接接受
        bossExecutor = Executors.newFixedThreadPool(1);
        // 从 Reactor：CPU核心数，处理IO读写
        int cpuCores = Runtime.getRuntime().availableProcessors();
        workerExecutor = Executors.newFixedThreadPool(cpuCores);
        // 业务线程池：2倍CPU核心数，处理业务逻辑
        businessExecutor = Executors.newFixedThreadPool(cpuCores * 2);
        
        // 初始化从 Reactor
        for (int i = 0; i < cpuCores; i++) {
            SubReactor subReactor = new SubReactor();
            subReactors.add(subReactor);
            workerExecutor.execute(subReactor);
        }
    }
    
    public void start(int port) throws IOException {
        ServerSocketChannel serverChannel = ServerSocketChannel.open();
        serverChannel.configureBlocking(false);
        serverChannel.socket().bind(new InetSocketAddress(port));
        
        // 主 Reactor 处理连接接受
        bossExecutor.execute(new BossReactor(serverChannel));
        
        System.out.println("多线程 NIO 服务器启动，监听端口 " + port);
        System.out.println("CPU 核心数: " + Runtime.getRuntime().availableProcessors());
        System.out.println("从 Reactor 数量: " + subReactors.size());
    }
    
    // 主 Reactor - 专门处理连接接受
    class BossReactor implements Runnable {
        private ServerSocketChannel serverChannel;
        private Selector selector;
        
        public BossReactor(ServerSocketChannel serverChannel) throws IOException {
            this.serverChannel = serverChannel;
            this.selector = Selector.open();
            serverChannel.register(selector, SelectionKey.OP_ACCEPT);
        }
        
        @Override
        public void run() {
            System.out.println("主 Reactor 启动...");
            try {
                while (!Thread.currentThread().isInterrupted()) {
                    selector.select();
                    Iterator<SelectionKey> keyIterator = selector.selectedKeys().iterator();
                    
                    while (keyIterator.hasNext()) {
                        SelectionKey key = keyIterator.next();
                        keyIterator.remove();
                        
                        if (key.isAcceptable()) {
                            handleAccept(key);
                        }
                    }
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        
        private void handleAccept(SelectionKey key) throws IOException {
            ServerSocketChannel serverChannel = (ServerSocketChannel) key.channel();
            SocketChannel clientChannel = serverChannel.accept();
            clientChannel.configureBlocking(false);
            
            // 轮询分配连接到从 Reactor
            SubReactor subReactor = getNextSubReactor();
            subReactor.registerClient(clientChannel);
            
            System.out.println("新客户端连接，分配到从 Reactor: " + 
                             subReactors.indexOf(subReactor));
        }
    }
    
    // 从 Reactor - 处理IO读写
    class SubReactor implements Runnable {
        private Selector selector;
        private BlockingQueue<SocketChannel> pendingChannels = new LinkedBlockingQueue<>();
        
        public SubReactor() throws IOException {
            this.selector = Selector.open();
        }
        
        public void registerClient(SocketChannel clientChannel) {
            // 将新连接添加到待处理队列
            pendingChannels.offer(clientChannel);
            // 唤醒 Selector
            selector.wakeup();
        }
        
        @Override
        public void run() {
            System.out.println("从 Reactor 启动: " + Thread.currentThread().getName());
            try {
                while (!Thread.currentThread().isInterrupted()) {
                    // 处理待注册的通道
                    processPendingChannels();
                    
                    // 等待IO事件
                    selector.select(1000); // 1秒超时，避免长时间阻塞
                    
                    if (selector.selectedKeys().isEmpty()) {
                        continue;
                    }
                    
                    Iterator<SelectionKey> keyIterator = selector.selectedKeys().iterator();
                    
                    while (keyIterator.hasNext()) {
                        SelectionKey key = keyIterator.next();
                        keyIterator.remove();
                        
                        try {
                            if (key.isReadable()) {
                                handleRead(key);
                            } else if (key.isWritable()) {
                                handleWrite(key);
                            }
                        } catch (IOException e) {
                            handleDisconnect(key);
                        }
                    }
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        
        private void processPendingChannels() throws IOException {
            SocketChannel clientChannel;
            while ((clientChannel = pendingChannels.poll()) != null) {
                clientChannel.register(selector, SelectionKey.OP_READ, new ClientSession());
                System.out.println("客户端注册到从 Reactor: " + 
                                 Thread.currentThread().getName());
            }
        }
        
        private void handleRead(SelectionKey key) throws IOException {
            SocketChannel clientChannel = (SocketChannel) key.channel();
            ClientSession session = (ClientSession) key.attachment();
            
            int bytesRead = clientChannel.read(session.readBuffer);
            
            if (bytesRead == -1) {
                handleDisconnect(key);
                return;
            }
            
            if (bytesRead > 0) {
                // 提交到业务线程池处理
                businessExecutor.execute(new BusinessProcessor(clientChannel, session));
            }
        }
        
        private void handleWrite(SelectionKey key) throws IOException {
            SocketChannel clientChannel = (SocketChannel) key.channel();
            ClientSession session = (ClientSession) key.attachment();
            
            synchronized (session.writeQueue) {
                while (!session.writeQueue.isEmpty()) {
                    ByteBuffer buffer = session.writeQueue.peek();
                    clientChannel.write(buffer);
                    
                    if (buffer.hasRemaining()) {
                        // 数据未写完，保持写兴趣
                        key.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
                        break;
                    } else {
                        // 数据已写完，移除队列
                        session.writeQueue.poll();
                        // 恢复只读兴趣
                        key.interestOps(SelectionKey.OP_READ);
                    }
                }
            }
        }
        
        private void handleDisconnect(SelectionKey key) {
            SocketChannel clientChannel = (SocketChannel) key.channel();
            try {
                System.out.println("客户端断开连接: " + clientChannel.getRemoteAddress());
                clientChannel.close();
            } catch (IOException e) {
                // 忽略
            }
            key.cancel();
        }
    }
    
    // 客户端会话
    class ClientSession {
        ByteBuffer readBuffer = ByteBuffer.allocate(1024);
        ByteBuffer lengthBuffer = ByteBuffer.allocate(4);
        boolean readingLength = true;
        int messageLength = 0;
        ByteBuffer messageBuffer;
        
        // 写队列，支持异步写入
        Queue<ByteBuffer> writeQueue = new ConcurrentLinkedQueue<>();
        
        void reset() {
            readingLength = true;
            messageLength = 0;
            lengthBuffer.clear();
            messageBuffer = null;
        }
    }
    
    // 业务处理器 - 在业务线程池中执行
    class BusinessProcessor implements Runnable {
        private SocketChannel clientChannel;
        private ClientSession session;
        
        public BusinessProcessor(SocketChannel clientChannel, ClientSession session) {
            this.clientChannel = clientChannel;
            this.session = session;
        }
        
        @Override
        public void run() {
            try {
                processReceivedData();
            } catch (Exception e) {
                System.err.println("业务处理错误: " + e.getMessage());
            }
        }
        
        private void processReceivedData() {
            session.readBuffer.flip();
            
            try {
                while (session.readBuffer.hasRemaining()) {
                    if (session.readingLength) {
                        int remaining = Math.min(session.readBuffer.remaining(), 
                                               session.lengthBuffer.remaining());
                        for (int i = 0; i < remaining; i++) {
                            session.lengthBuffer.put(session.readBuffer.get());
                        }
                        
                        if (!session.lengthBuffer.hasRemaining()) {
                            session.lengthBuffer.flip();
                            session.messageLength = session.lengthBuffer.getInt();
                            
                            if (session.messageLength <= 0 || session.messageLength > 1024 * 1024) {
                                session.reset();
                                continue;
                            }
                            
                            session.messageBuffer = ByteBuffer.allocate(session.messageLength);
                            session.readingLength = false;
                            session.lengthBuffer.clear();
                        }
                    } else {
                        int remaining = Math.min(session.readBuffer.remaining(), 
                                               session.messageBuffer.remaining());
                        for (int i = 0; i < remaining; i++) {
                            session.messageBuffer.put(session.readBuffer.get());
                        }
                        
                        if (!session.messageBuffer.hasRemaining()) {
                            session.messageBuffer.flip();
                            String fullMessage = charset.decode(session.messageBuffer).toString();
                            
                            // 模拟业务处理耗时
                            processBusinessLogic(fullMessage);
                            
                            // 发送响应
                            sendResponse("已处理: " + fullMessage.length() + " 字符");
                            
                            session.reset();
                        }
                    }
                }
            } finally {
                session.readBuffer.compact();
            }
        }
        
        private void processBusinessLogic(String message) {
            // 模拟业务处理耗时
            try {
                Thread.sleep(10); // 10ms 业务处理时间
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
            
            System.out.println("[" + Thread.currentThread().getName() + 
                             "] 处理消息: " + message.length() + " 字符");
        }
        
        private void sendResponse(String response) {
            try {
                byte[] responseData = response.getBytes(charset);
                ByteBuffer responseBuffer = ByteBuffer.allocate(4 + responseData.length);
                
                responseBuffer.putInt(responseData.length);
                responseBuffer.put(responseData);
                responseBuffer.flip();
                
                // 添加到写队列
                synchronized (session.writeQueue) {
                    session.writeQueue.offer(responseBuffer);
                }
                
                // 唤醒对应的从 Reactor 进行写入
                for (SubReactor subReactor : subReactors) {
                    if (subReactor.selector.keys().contains(clientChannel.keyFor(subReactor.selector))) {
                        SelectionKey key = clientChannel.keyFor(subReactor.selector);
                        key.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
                        subReactor.selector.wakeup();
                        break;
                    }
                }
                
            } catch (Exception e) {
                System.err.println("发送响应失败: " + e.getMessage());
            }
        }
    }
    
    private SubReactor getNextSubReactor() {
        // 简单的轮询负载均衡
        SubReactor reactor = subReactors.get(subReactorIndex);
        subReactorIndex = (subReactorIndex + 1) % subReactors.size();
        return reactor;
    }
    
    public void shutdown() {
        bossExecutor.shutdown();
        workerExecutor.shutdown();
        businessExecutor.shutdown();
        
        try {
            if (!bossExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                bossExecutor.shutdownNow();
            }
            if (!workerExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                workerExecutor.shutdownNow();
            }
            if (!businessExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                businessExecutor.shutdownNow();
            }
        } catch (InterruptedException e) {
            bossExecutor.shutdownNow();
            workerExecutor.shutdownNow();
            businessExecutor.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }
    
    public static void main(String[] args) throws IOException {
        MultiThreadNioServer server = new MultiThreadNioServer();
        try {
            server.start(8888);
            
            // 添加关闭钩子
            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                System.out.println("正在关闭服务器...");
                server.shutdown();
            }));
            
        } catch (IOException e) {
            e.printStackTrace();
            server.shutdown();
        }
    }
}